import abstract_decoder


class CaesarDecoder(abstract_decoder.AbstractDecoder):
    """
    凯撒密码解密，密码内字母根据偏移量进行移位产生密文，类似的方法进行解析
    当前类根据CTF常用的flag或ctf标志来作为参考，自动计算偏移量，然后解密
    """

    def __init__(self):
        self.shift = None
        super().__init__("Caesar")

    def find_key(self, data:str):
        """
        自动计算偏移量
        """
        keywords = ["flag", "ctf"]
        for i in range(26):
            for keyword in keywords:
                tmp = self.encode(keyword, i)
                if data.find(tmp) >= 0:
                    self.shift = i
                    print(f"caesar found key: {self.shift}")
                    return True
        return False

    def encode(self, source, shift_num):
        """
        简易加密
        """
        result = ""
        for i in range(len(source)):
            if self.is_valid_letter(source[i]):
                result += self.shift_char(source[i], shift_num)
            else:
                result += source[i]
        return result

    def decode(self, data, key=None):
        """
        解密
        """
        if self.shift is None:
            self.find_key(data)

        result = ""
        for char in data:
            if self.is_valid_letter(char):
                result += self.shift_char(char, -self.shift)
            else:
                result += char
        return result

    def is_valid_letter(self, char):
        """
        判断是否为合法英文字符
        """
        return 65 <= ord(char) <= 90 or 97 <= ord(char) <= 122
    
    def shift_char(self, char, shift_num):
        """
        加密单个字符
        """
        if self.is_valid_letter(char):
            if 65 <= ord(char) <= 90:
                return chr((ord(char) - 65 + shift_num) % 26 + 65)
            else:
                return chr((ord(char) - 97 + shift_num) % 26 + 97)
        else:
            return char

    def is_valid(self, data, key=None):
        self.find_key(data)
        if self.shift is None or self.shift == 0:
            return False
        return True
    
    def reset(self):
        self.shift = None

if __name__ == "__main__":
    decoder = CaesarDecoder()    
    with open("test-data/caesar.txt", "r") as f:
        data = f.read()
        print(decoder.is_valid(data))
        print(decoder.decode(data))
